autodiff rlib handling #149033

ZuseZ4 · 2025-11-17T17:41:40Z

As I learned recently, we now apparently support rlib builds already in some cases.
With the last hint from saethlin this seems to now cover all cases. To be sure I'll add a few more testcases before I mark it as ready.

Once this PR lands, we should to the best of my knowledge, support autodiff in almost code locations, only vtable/dyn ptr remain unsupported for now.

r? ghost

rust-log-analyzer · 2025-11-17T17:47:06Z

The job tidy failed! Check out the build log: (web) (plain enhanced) (plain)

Click to see the possible cause of the failure (guessed by this bot)

[TIMING:end] tool::ToolBuild { build_compiler: Compiler { stage: 0, host: x86_64-unknown-linux-gnu, forced_compiler: false }, target: x86_64-unknown-linux-gnu, tool: "tidy", path: "src/tools/tidy", mode: ToolBootstrap, source_type: InTree, extra_features: [], allow_features: "", cargo_args: [], artifact_kind: Binary } -- 11.350
[TIMING:end] tool::Tidy { compiler: Compiler { stage: 0, host: x86_64-unknown-linux-gnu, forced_compiler: false }, target: x86_64-unknown-linux-gnu } -- 0.000
fmt check
Diff in /checkout/compiler/rustc_codegen_llvm/src/intrinsic.rs:15:
 use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv, LayoutOf};
 use rustc_middle::ty::{self, GenericArgsRef, Instance, SimdAlign, Ty, TyCtxt, TypingEnv};
 use rustc_middle::{bug, span_bug};
+use rustc_session::config::CrateType;
 use rustc_span::{Span, Symbol, sym};
 use rustc_symbol_mangling::{mangle_internal_symbol, symbol_name_for_instance_in_crate};
 use rustc_target::callconv::PassMode;
Diff in /checkout/compiler/rustc_codegen_llvm/src/intrinsic.rs:21:
 use rustc_target::spec::Os;
 use tracing::debug;
 
-use rustc_session::config::CrateType;
-
 use crate::abi::FnAbiLlvmExt;
 use crate::builder::Builder;
 use crate::builder::autodiff::{adjust_activity_to_abi, generate_enzyme_call};
Diff in /checkout/compiler/rustc_codegen_llvm/src/intrinsic.rs:1146:
             let _ = tcx.dcx().emit_almost_fatal(AutoDiffWithoutLto);
         }
     } else {
-        if lto != rustc_session::config::Lto::Fat && !tcx.sess.opts.cg.linker_plugin_lto.enabled(){
+        if lto != rustc_session::config::Lto::Fat && !tcx.sess.opts.cg.linker_plugin_lto.enabled() {
             let _ = tcx.dcx().emit_almost_fatal(AutoDiffWithoutLto);
         }
     }
Diff in /checkout/compiler/rustc_mir_transform/src/cross_crate_inline.rs:34:
         return true;
     }
 
-    if tcx.has_attr(def_id, sym::autodiff_forward) || tcx.has_attr(def_id, sym::autodiff_reverse) || tcx.has_attr(def_id, sym::rustc_autodiff) {
+    if tcx.has_attr(def_id, sym::autodiff_forward)
+        || tcx.has_attr(def_id, sym::autodiff_reverse)
+        || tcx.has_attr(def_id, sym::rustc_autodiff)
+    {
         return true;
     }
 
fmt: checked 6549 files

bjorn3 · 2025-11-17T19:56:55Z

compiler/rustc_mir_transform/src/cross_crate_inline.rs


+    if tcx.has_attr(def_id, sym::autodiff_forward) || tcx.has_attr(def_id, sym::autodiff_reverse) || tcx.has_attr(def_id, sym::rustc_autodiff) {
+        return true;
+    }


Is rustc_monomorphize::collector::autodiff::collect_autodiff_fn even necessary in the first place. The catch_unwind intrinsic also needs the function arguments to be codegened, yet it doesn't have any special casing in the monomorphization collector. Enzyme needs the function to differentiate to be in the same module, but I would expect fat LTO to run early enough to already cause that to happen. If you remove collect_autodiff_fn then I would assume this special case in cross_crate_inlineable isn't necessary either.

@bjorn3 Unfortunately it seems to be. I've tried removing both and we are back to an old bug, could not find source function.
In the beginning we had this case when we'd differentiate f to generate g, but then only call g, not f.
During the lowering we'd realize that f is unused, delete it, and then fail on the llvm-ir level, since we now can't generate the body of g since it depended on f. Our new intrinsic in g in theory references the source function f, but it seems that the monomorphize adjustments are still needed.

error: internal compiler error: compiler/rustc_codegen_llvm/src/intrinsic.rs:1183:9: could not find source function thread 'rustc' (815281) panicked at compiler/rustc_codegen_llvm/src/intrinsic.rs:1183:9: Box<dyn Any> stack backtrace: 0: 0x72cccc49e5c2 - <std::sys::backtrace::BacktraceLock::print::DisplayBacktrace as core::fmt::Display>::fmt::h47627beda0767f4c 1: 0x72cccc4cc9af - core::fmt::write::h8a63064c67176fd8 2: 0x72cccc4516e3 - std::io::Write::write_fmt::h5e9b646cef91150b 3: 0x72cccc466ae2 - std::sys::backtrace::BacktraceLock::print::h4634d27c0afbaebb 4: 0x72cccc46ee7a - std::panicking::default_hook::{{closure}}::h23e0299d4738b766 5: 0x72cccc46ecd4 - std::panicking::default_hook::hfeec6294e8751b62 6: 0x72ccc8bd3a36 - std[e20d0b87aa1ec278]::panicking::update_hook::<alloc[d35db71d034d3387]::boxed::Box<rustc_driver_impl[ee6e3d57fbee4319]::install_ice_hook::{closure#1}>>::{closure#0} 7: 0x72cccc46f52f - std::panicking::panic_with_hook::hd57cd463d793ed79 8: 0x72cccbf6b8c3 - std[e20d0b87aa1ec278]::panicking::begin_panic::<rustc_errors[52538de81f6f8989]::ExplicitBug>::{closure#0} 9: 0x72cccbf6a976 - std[e20d0b87aa1ec278]::sys::backtrace::__rust_end_short_backtrace::<std[e20d0b87aa1ec278]::panicking::begin_panic<rustc_errors[52538de81f6f8989]::ExplicitBug>::{closure#0}, !> 10: 0x72cccbf63f5a - std[e20d0b87aa1ec278]::panicking::begin_panic::<rustc_errors[52538de81f6f8989]::ExplicitBug> 11: 0x72cccbf4c652 - <rustc_errors[52538de81f6f8989]::diagnostic::BugAbort as rustc_errors[52538de81f6f8989]::diagnostic::EmissionGuarantee>::emit_producing_guarantee 12: 0x72cccbce40c5 - rustc_middle[4d11da33ed8b74e2]::util::bug::opt_span_bug_fmt::<rustc_span[eee2686a21fb97e]::span_encoding::Span>::{closure#0} 13: 0x72cccbce413a - rustc_middle[4d11da33ed8b74e2]::ty::context::tls::with_opt::<rustc_middle[4d11da33ed8b74e2]::util::bug::opt_span_bug_fmt<rustc_span[eee2686a21fb97e]::span_encoding::Span>::{closure#0}, !>::{closure#0} 14: 0x72cccbce083b - rustc_middle[4d11da33ed8b74e2]::ty::context::tls::with_context_opt::<rustc_middle[4d11da33ed8b74e2]::ty::context::tls::with_opt<rustc_middle[4d11da33ed8b74e2]::util::bug::opt_span_bug_fmt<rustc_span[eee2686a21fb97e]::span_encoding::Span>::{closure#0}, !>::{closure#0}, !> 15: 0x72cccbceabe2 - rustc_middle[4d11da33ed8b74e2]::util::bug::bug_fmt 16: 0x72ccc8f90449 - rustc_codegen_llvm[fd0152a31030d463]::intrinsic::codegen_autodiff 17: 0x72ccc90aa0fe - <rustc_codegen_llvm[fd0152a31030d463]::builder::GenericBuilder<rustc_codegen_llvm[fd0152a31030d463]::context::FullCx> as rustc_codegen_ssa[78843d74ba320aa6]::traits::intrinsic::IntrinsicCallBuilderMethods>::codegen_intrinsic_call 18: 0x72ccc903e8b5 - <rustc_codegen_ssa[78843d74ba320aa6]::mir::FunctionCx<rustc_codegen_llvm[fd0152a31030d463]::builder::GenericBuilder<rustc_codegen_llvm[fd0152a31030d463]::context::FullCx>>>::codegen_intrinsic_call 19: 0x72ccc9033c07 - rustc_codegen_ssa[78843d74ba320aa6]::mir::codegen_mir::<rustc_codegen_llvm[fd0152a31030d463]::builder::GenericBuilder<rustc_codegen_llvm[fd0152a31030d463]::context::FullCx>> 20: 0x72ccc90cb8c1 - rustc_codegen_ssa[78843d74ba320aa6]::base::codegen_instance::<rustc_codegen_llvm[fd0152a31030d463]::builder::GenericBuilder<rustc_codegen_llvm[fd0152a31030d463]::context::FullCx>> 21: 0x72ccc90bda96 - <rustc_middle[4d11da33ed8b74e2]::mir::mono::MonoItem as rustc_codegen_ssa[78843d74ba320aa6]::mono_item::MonoItemExt>::define::<rustc_codegen_llvm[fd0152a31030d463]::builder::GenericBuilder<rustc_codegen_llvm[fd0152a31030d463]::context::FullCx>> 22: 0x72ccc90921c1 - rustc_codegen_llvm[fd0152a31030d463]::base::compile_codegen_unit::module_codegen 23: 0x72ccc9082313 - rustc_codegen_llvm[fd0152a31030d463]::base::compile_codegen_unit 24: 0x72ccc90cadec - rustc_codegen_ssa[78843d74ba320aa6]::base::codegen_crate::<rustc_codegen_llvm[fd0152a31030d463]::LlvmCodegenBackend> 25: 0x72ccc8f630c6 - <rustc_codegen_llvm[fd0152a31030d463]::LlvmCodegenBackend as rustc_codegen_ssa[78843d74ba320aa6]::traits::backend::CodegenBackend>::codegen_crate 26: 0x72ccc8ebf84e - <rustc_session[85ae29ebbb651820]::session::Session>::time::<alloc[d35db71d034d3387]::boxed::Box<dyn core[5b293a957d0a6e5f]::any::Any>, rustc_interface[80f9411e4dfda4b]::passes::start_codegen::{closure#0}> 27: 0x72ccc8dee328 - rustc_interface[80f9411e4dfda4b]::passes::start_codegen 28: 0x72ccc8ed10ba - <rustc_interface[80f9411e4dfda4b]::queries::Linker>::codegen_and_build_linker

Right. Maybe the reason why it works for catch_unwind is that catch_unwind accepts function pointers rather than function items as arguments?

Is it possible to replace fetching #[rustc_autodiff] from the function to be differentiated with a regular argument (or const generic) passed to the autodiff intrinsic, pass a function pointer to the autodiff intrinsic and then get the symbol name from the function pointer value (under the assumption that the function pointer value is a compile-time known value, which should hold true for the #[autodiff_forward] and #[autodiff_backward] expansions) without involving Instance?

Edit: I got the source and differentiated function backwards. #[rustc_autodiff] is applied to the differentiated function which contains the autodiff call, so that one would already work. You would only need to replace the function item with a function pointer and adjust adjust_activity_to_abi to accept the function pointer type rather than a concrete Instance and then use fn_abi_of_fn_ptr instead of fn_abi_of_instance.

That might be a reason, though it feels weird to me that a fn item does not count as use, whereas a function pointer does. And yes, they should always be compile-time known for now.
We have to add support for dyn (vtable) in the future, though, which in other languages works by differentiating all candidates, and then at runtime looking up which primal function was given and picking the responding differentiated function. That might require extra work anyway, so I'm also fine with not considering it for now.

For reference, here is the current expansion of placing #[autodiff_forward(...)] on a function _f1.

#[rustc_autodiff] pub fn _f1(x: f64) -> f64 { simple_dep::f(x, x) * f2(x) } #[rustc_autodiff(Forward, 1, Dual, Dual)] pub fn df1_lib(x: f64, bx_0: f64) -> (f64, f64) { ::core::intrinsics::autodiff(_f1, df1_lib, (x, bx_0)) }

cc @oli-obk for thoughts and @Sa4dUs who developed our new intrinsic. I think moving it shouldn't be too much work?

it shouldn't, basically replacing the current logic in cg_llvm::intrinsic::codegen_autodiff that starts from the DefId and generic args to obtain the value

rlib handling

909a480

ZuseZ4 added the F-autodiff `#![feature(autodiff)]` label Nov 17, 2025

bjorn3 reviewed Nov 17, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

autodiff rlib handling #149033

autodiff rlib handling #149033

ZuseZ4 commented Nov 17, 2025

Uh oh!

rust-log-analyzer commented Nov 17, 2025

Uh oh!

bjorn3 Nov 17, 2025

Uh oh!

ZuseZ4 Nov 17, 2025 •

edited

Loading

Uh oh!

bjorn3 Nov 17, 2025 •

edited

Loading

Uh oh!

ZuseZ4 Nov 17, 2025

Uh oh!

Sa4dUs Nov 18, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

autodiff rlib handling #149033

Are you sure you want to change the base?

autodiff rlib handling #149033

Conversation

ZuseZ4 commented Nov 17, 2025

Uh oh!

rust-log-analyzer commented Nov 17, 2025

Uh oh!

bjorn3 Nov 17, 2025

Choose a reason for hiding this comment

Uh oh!

ZuseZ4 Nov 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bjorn3 Nov 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ZuseZ4 Nov 17, 2025

Choose a reason for hiding this comment

Uh oh!

Sa4dUs Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ZuseZ4 Nov 17, 2025 •

edited

Loading

bjorn3 Nov 17, 2025 •

edited

Loading